package websocket

import (
	"fmt"
	"github.com/gorilla/websocket"
	"net/http"
	"sync"
)

var (
	defaultMaxConnNum int64 = 10240
	defaultMaxMsgLen  int64 = 10240 * 128
)

// User 连接信息
type User struct {
	conn *websocket.Conn
	msg  chan []byte
}

// WsServer websocket连接对象
type WsServer struct {
	maxConnNum int64               // 最大连接数
	maxMsgLen  int64               // 消息最大长度
	users      map[*User]bool      // 用户列表
	upgrader   *websocket.Upgrader // 连接升级 http --> websocket
	register   chan *User          // 注册用户
	unRegister chan *User          // 注销用户
	broadcast  chan []byte         // 广播消息
	mtx        *sync.Mutex
}

// newUpgrader 创建连接升级器
func newUpgrader() *websocket.Upgrader {
	return &websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			if r.Method != http.MethodGet {
				return false
			}
			// todo 关于连接验证在这里做
			return true
		},
	}
}

// NewWsServer 创建WebSocket服务
func NewWsServer(maxConnNum, maxMsgLen int64) *WsServer {
	return &WsServer{
		maxConnNum: maxConnNum,
		maxMsgLen:  maxMsgLen,
		users:      make(map[*User]bool),
		upgrader:   newUpgrader(),
		register:   make(chan *User),
		unRegister: make(chan *User),
		broadcast:  make(chan []byte),
		mtx:        &sync.Mutex{},
	}
}

// NewDefaultWsServer 创建默认WS服务
func NewDefaultWsServer() *WsServer {
	return NewWsServer(defaultMaxConnNum, defaultMaxMsgLen)
}

// Handel 连接处理函数
func (ws *WsServer) Handel(w http.ResponseWriter, r *http.Request) {
	conn, err := ws.upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Printf("http to websocket err:%v\n", err)
		return
	}
	conn.SetReadLimit(ws.maxMsgLen)

	ws.mtx.Lock()
	if int64(len(ws.users)) >= ws.maxConnNum {
		ws.mtx.Unlock()
		_ = conn.Close()
		fmt.Printf("this conn more than maxConnNum")
		return
	}
	// 用户注册
	user := &User{
		conn: conn,
		msg:  make(chan []byte),
	}
	ws.register <- user
	// 最后注销
	defer func() {
		ws.unRegister <- user
	}()

	// 处理数据
	go ws.read(user)
	ws.write(user)

}

// read 读取连接信息
func (ws *WsServer) read(user *User) {
	// 循环读取消息
	for {
		_, message, err := user.conn.ReadMessage()
		if err != nil {
			ws.mtx.Lock()
			fmt.Printf("address:%s user quit\n", user.conn.RemoteAddr().String())
			ws.unRegister <- user
			break
		}
		// 把消息发送到广播消息中
		ws.broadcast <- message
	}
}

// write 写入消息
func (ws *WsServer) write(user *User) {
	for data := range user.msg {
		err := user.conn.WriteMessage(websocket.TextMessage, data)
		if err != nil {
			fmt.Printf("write err:%v\n", err)
			break
		}
	}
}

// DealHandel 处理消息
func (ws *WsServer) DealHandel() {
	for {
		select {
		// 处理注册消息
		case user := <-ws.register:
			// 将用户添加到用户管理列表中
			ws.users[user] = true
			ws.mtx.Unlock()
		// 处理注销消息
		case user := <-ws.unRegister:
			// 移除用户管理器用户
			if _, ok := ws.users[user]; ok {
				_ = user.conn.Close()
				delete(ws.users, user)
				ws.mtx.Unlock()
			}
		// 处理广播消息
		case data := <-ws.broadcast:
			// 将消息广播给每一个用户
			for user := range ws.users {
				select {
				case user.msg <- data:
				default:
					ws.mtx.Lock()
					delete(ws.users, user)
					close(user.msg)
					ws.mtx.Unlock()
				}
			}
		}
	}
}

// Close 关闭服务
func (ws *WsServer) Close() {
	ws.mtx.Lock()
	defer ws.mtx.Unlock()
	for user := range ws.users {
		_ = user.conn.Close()
	}
	ws.users = nil
}
